#   This program is free software; you can redistribute it and/or modify
#   it under the terms of the version 3 of the GNU Lesser General Public License
#   as published by the Free Software Foundation.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU Lesser General Public License
#   along with this program; if not, write to the Free Software
#   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
#
# Copyright (c) NEC Deutschland GmbH, NEC HPC Europe
#
# $Id$

from time import sleep
from json import dumps as json_dumps, loads as json_loads 
import amqplib.client_0_8 as amqp
import re

from basic_types.marshaller import PrimitiveTypesMarshaller
from basic_types.marshaller import PrimitiveTypesUnmarshaller

from aggmon.channel import Channel
from aggmon.stoppable_thread import StoppableThread
from basic_types.factory import Factory


__all__ = ["AMQPChannel"]


class AMQPCallbackThread(StoppableThread):
    def __init__( self, amqp_channel, topic, notify_callback, raw=False ):
        StoppableThread.__init__( self )
        self.__notify_callback = notify_callback
        self.__parent_channel = amqp_channel
        self.__topic = topic
        self.cb_channel = amqp_channel.get_connection().channel()
        self.cb_channel.access_request( amqp_channel.channel_name, active=True, read=True )

        queue_name, _messages, _consumers = self.cb_channel.queue_declare( exclusive=True, auto_delete=True )
        self.queue_name = queue_name
        self.cb_channel.queue_bind( queue_name, self.__parent_channel.get_exchange_name(), routing_key=topic )
        if raw:
            self.cb_channel.basic_consume( queue_name, callback=notify_callback, consumer_tag=queue_name,
                                           no_ack=False )
        else:
            self.cb_channel.basic_consume( queue_name, callback=self.__callback, consumer_tag=queue_name,
                                           no_ack=False )
        self.start()

    def __callback( self, amqp_message ):
        amqp_message.channel.basic_ack( amqp_message.delivery_tag )

        # ignore special message that helps terminating the thread
        if amqp_message.body == "quit":
            return

        jsonrpc_request = json_loads( amqp_message.body )
        message = Factory.build( jsonrpc_request["method"] )
        message.unmarshal( PrimitiveTypesUnmarshaller( jsonrpc_request["params"] ) )
        topic = amqp_message.routing_key
        #print "unmarshalled message:", message
        self.__notify_callback( message, topic=topic )

    def join( self ):
        # send a message to kick the wait()
        amqp_message = amqp.Message( "quit", content_type="application/timacs-metric", application_headers={} )
        self.__parent_channel.get_channel().basic_publish( amqp_message, self.__parent_channel.get_exchange_name(), routing_key=self.__topic )
        StoppableThread.join(self)

    # Thread
    def run( self ):            
        while not self._stopevent.isSet():
            try:
                self.cb_channel.wait()
            except Exception, e:
                print "Exception in subscriber!"
                print e
                raise
        self.cb_channel.basic_cancel( self.queue_name )
        #self.cb_channel.queue_delete( self.queue_name )


class AMQPChannel(Channel):
    def __init__( self, url, durable=False ):
        Channel.__init__( self )
        self.url = url

        scheme, rest = url.split( "://" )
        assert scheme == "amqp"
        netloc, channel_name = rest.split( "/" )
        userid_password, self.__host = netloc.split( "@" )
        userid, password = userid_password.split( ":" )
        self.channel_name = "/data/" + channel_name 

        self.__amqp_exchange_name = channel_name

        self.__amqp_connection = amqp.Connection( self.__host, userid=userid, password=password )
        self.__amqp_channel = self.__amqp_connection.channel()
        self.__amqp_channel.access_request( self.channel_name, active=True, write=True )
        self.__amqp_channel.exchange_declare( self.__amqp_exchange_name, "topic", auto_delete=False, durable=durable )
        self.cb_threads = []
        # testing serialized handling of subscriber callbacks because pyamqp is soooo bad
        self.__subscriber_notify_callbacks = []

    def cleanup( self ):
        # stop subscriber threads
        for thread in self.cb_threads:
            thread.join()
        self.__amqp_channel.close()
        self.__amqp_connection.close()

    def get_channel( self ):
        return self.__amqp_channel

    def get_connection( self ):
        return self.__amqp_connection

    def get_exchange_name( self ):
        return self.__amqp_exchange_name

    # Channel
    def publish( self, topic, message, application_headers={}, content_type="application/timacs-metric", raw=False, **kwds ):
        """
        Publish a message to the AMQP exchange the channel is pointing at.
        When the raw flag is True: interpret the message as payload. Otherwise encapsulate it into a JSON
        request, passing 'method', 'params' and 'id'.
        AMQP message properties such as content_type, app_id, type, reply_to, etc... can be passed along
        with the publish() call. They will be inserted into the AMQP message.
        """
        if raw:
            payload = message
        else:
            params = {}
            message.marshal( PrimitiveTypesMarshaller( params ) )
            jsonrpc_request = { "method": message.__class__.__name__.rsplit( "." )[-1], "params": params, "id": None }
            payload = json_dumps( jsonrpc_request )
            #print "publishing:", payload
        amqp_message = amqp.Message( payload, content_type=content_type, application_headers=application_headers, **kwds )
        try:
            self.__amqp_channel.basic_publish( amqp_message, self.__amqp_exchange_name, routing_key=topic )
        except Exception, e:
            print "Exception in publish!"
            print e
            raise

    def _serializedSubscriberWorker( self, message, topic=None ):
        for (subscriber_topic, subscriber_callback) in self.__subscriber_notify_callbacks:
            pattern = re.sub( r'\.', "\.", subscriber_topic )
            pattern = re.sub( r'\*', "[^\.]*", pattern )
            pattern = re.sub( r'#', ".*", pattern )
            if re.match( pattern, topic ) is not None:
            #    #print "Matched %s to topic %s\n" % (pattern, topic)
                subscriber_callback( message, topic=topic )


    def subscribe( self, topic, notify_callback, raw=False, emulate_topic=True ):
        """
        Subscribe to a topic, execute notify_callback on every event received.
        When the raw flag is set to True the callback gets passed the raw AMQP messages
        instead of JSON deserialized stuff.
        """
        if not emulate_topic or raw:
            # WARNING: there is an issue with pyamqp and multiple subscribers, probably due
            # to the fact that pyamqp is not thread-safe. Use raw only if you have only one subscriber
            # on the channel/connection!
            self.cb_threads.append( AMQPCallbackThread( self, topic, notify_callback, raw=raw ) )
        else:
            if len( self.cb_threads ) == 0:
                self.cb_threads.append( AMQPCallbackThread( self, '#', self._serializedSubscriberWorker, raw=False ) )
            self.__subscriber_notify_callbacks.append( (topic, notify_callback) )
